import numpy as np
import pandas as pd
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

data_path = r'LinearUse.csv'
data = pd.read_csv(data_path)

# 选择降水量列 (假设降雨量在第三列
rainfall = data['rrr'].values

# 数据归一化
scaler = MinMaxScaler(feature_range=(0, 1))
rainfall = scaler.fit_transform(rainfall.reshape(-1, 1))


# 创建数据集
def create_dataset(data, time_step=1):
    X, y = [], []
    for i in range(len(data) - time_step - 1):
        a = data[i:(i + time_step), 0]
        X.append(a)
        y.append(data[i + time_step, 0])
    return np.array(X), np.array(y)


time_step = 48  # 使用过去4年的数据来预测下一个月的降雨量
X, y = create_dataset(rainfall, time_step)
X = X.reshape(X.shape[0], X.shape[1], 1)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 构建LSTM模型
model = Sequential()
model.add(LSTM(50, return_sequences=True, input_shape=(time_step, 1)))
model.add(LSTM(50))
model.add(Dense(1))

model.compile(optimizer='adam', loss='mean_squared_error')

# 训练模型
model.fit(X_train, y_train, epochs=500, batch_size=32, verbose=1)

# 预测第十年的月均降雨量
train_predict = model.predict(X_train)
test_predict = model.predict(X_test)
# 反归一化预测结果
train_predict = scaler.inverse_transform(train_predict)
y_train = scaler.inverse_transform([y_train])
test_predict = scaler.inverse_transform(test_predict)
y_test = scaler.inverse_transform([y_test])

# 绘制预测结果和真实值的图表
plt.figure(figsize=(15, 6))
plt.plot(y_train[0], label='Real Train Data')
plt.plot(train_predict[:, 0], label='Train Prediction')
plt.plot(y_test[0], label='Real Test Data')
plt.plot(test_predict[:, 0], label='Test Prediction')
plt.title('LSTM Rainfall Prediction')
plt.xlabel('Time Step')
plt.ylabel('Rainfall')
plt.legend()
plt.show()